import sys
sys.path.append('../')

import os
import pandas as pd
import numpy as np
from logparser import Spell, Drain
from tqdm import tqdm
from logdeep.dataset.session import sliding_window
import torch
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity, euclidean_distances

tqdm.pandas()
pd.options.mode.chained_assignment = None  # default='warn'
process_evolution = True

def find_most_similar_eventId(df_templates, df_templates_evo):
    # 构造存储 eventId 对应关系的字典
    eventId_mapping = {}
    
    # 将 Vector 列转为 numpy 数组以便计算相似度
    df_templates['Vector'] = df_templates['Vector'].apply(np.array)
    df_templates_evo['Vector'] = df_templates_evo['Vector'].apply(np.array)
    print(type(df_templates_evo['Vector'].iloc[0]))
    
    for index, row in df_templates_evo.iterrows():
        evo_eventId = row['EventId']
        evo_vector = row['Vector']
        
        # 检查 eventId 是否在 df_templates 中
        if evo_eventId not in df_templates['EventId'].values:
            # 计算余弦相似度和欧几里得距离
            cosine_similarities = cosine_similarity([evo_vector], df_templates['Vector'].tolist()).flatten()
            distances = euclidean_distances([evo_vector], df_templates['Vector'].tolist()).flatten()
            combined_similarity = distances - cosine_similarities
            
            # 找到相似度最高的模板对应的 eventId
            max_similarity_index = np.argmin(combined_similarity)
            most_similar_eventId = df_templates.iloc[max_similarity_index]['EventId']

            # 构造 eventId 对应关系
            eventId_mapping[evo_eventId] = most_similar_eventId
    for key in eventId_mapping:
        if key != eventId_mapping[key]:
            print("eventId_mapping: ",key, eventId_mapping[key])
            
    return eventId_mapping

# In the first column of the log, "-" indicates non-alert messages while others are alert messages.
def count_anomaly(log_path):
    total_size = 0
    normal_size = 0
    with open(log_path, errors='ignore') as f:
        for line in f:
            total_size += 1
            if line.split('')[0] == '-':
                normal_size += 1
    print("total size {}, abnormal size {}".format(total_size, total_size - normal_size))


def deeplog_file_generator(filename, df, features):
    with open(filename, 'w') as f:
        for _, row in df.iterrows():
            for val in zip(*row[features]):
                f.write(','.join([str(v) for v in val]) + ' ')
            f.write('\n')


def parse_log(input_dir, output_dir, log_file, parser_type):
    log_format = '<Label> <Id> <Date> <Admin> <Month> <Day> <Time> <AdminAddr> <Content>'
    regex = [
        r'(0x)[0-9a-fA-F]+',  # hexadecimal
        r'\d+\.\d+\.\d+\.\d+',
        r'(?<=Warning: we failed to resolve data source name )[\w\s]+',
        r'\d+'
    ]
    keep_para = False
    if parser_type == "drain":
        # the hyper parameter is set according to http://jmzhu.logpai.com/pub/pjhe_icws2017.pdf
        st = 0.3  # Similarity threshold
        depth = 3  # Depth of all leaf nodes

        # Drain is modified
        parser = Drain.LogParser(log_format,
                                 indir=input_dir,
                                 outdir=output_dir,
                                 depth=depth,
                                 st=st,
                                 rex=regex,
                                 keep_para=keep_para, maxChild=1000)
        parser.parse(log_file)

    elif parser_type == "spell":
        tau = 0.35
        parser = Spell.LogParser(indir=data_dir,
                                 outdir=output_dir,
                                 log_format=log_format,
                                 tau=tau,
                                 rex=regex,
                                 keep_para=keep_para)
        parser.parse(log_file)



def sample_raw_data(data_file, output_file, sample_window_size, sample_step_size):
    # sample 1M by sliding window, abnormal rate is over 2%
    sample_data = []
    labels = []
    idx = 0

    # spirit dataset can start from the 2Mth line, as there are many abnormal lines gathering in the first 2M
    with open(data_file, 'r', errors='ignore') as f:
        for line in f:
            labels.append(line.split()[0] != '-')
            sample_data.append(line)

            if len(labels) == sample_window_size:
                abnormal_rate = sum(np.array(labels)) / len(labels)
                print(f"{idx + 1} lines, abnormal rate {abnormal_rate}")
                break

            idx += 1
            if idx % sample_step_size == 0:
                print(f"Process {round(idx/sample_window_size * 100,4)} % raw data", end='\r')

    with open(output_file, "w") as f:
        f.writelines(sample_data)

    print("Sampling done")


if __name__ == "__main__":
    data_dir = "../../logbert/dataset/tbird/"
    output_dir = "../output/tbird/"
    raw_log_file = "Thunderbird.log"
    sample_log_file = "Thunderbird_20M.log"
    sample_window_size = 2*10**7
    sample_step_size = 10**4
    window_name = ''
    log_file = sample_log_file

    parser_type = 'drain'
    #mins
    window_size = 1
    step_size = 0.5
    train_ratio = 0.4

    ########
    # count anomaly
    ########
    # count_anomaly(data_dir + log_file)
    # sys.exit()

    #########
    # sample raw data
    #########
    sample_raw_data(data_dir+raw_log_file, data_dir+sample_log_file, sample_window_size, sample_step_size )


    ##########
    # Parser #
    #########
    # parse_log(data_dir, output_dir, log_file, parser_type)

    ##################
    # Transformation #
    ##################
    if process_evolution == False:
        df = pd.read_csv(f'{output_dir}{log_file}_structured.csv')
    else:
        df = pd.read_csv(f'{output_dir}{log_file}_structured_evolution.csv')
        # 抽取演化后日志模板的语义向量
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2', device=device)

        df_templates_evo = pd.read_csv(f'{output_dir}{log_file}_templates_evolution.csv')
        embeddings = model.encode(df_templates_evo['EventTemplate'].tolist())  # num_workers=num_workers)
        df_templates_evo['Vector'] = [embedding.tolist() for embedding in embeddings]
        template_dict = df_templates_evo.set_index('EventTemplate')['Vector'].to_dict()

        df_templates = pd.read_csv(f'{output_dir}{log_file}_templates.csv')
        embeddings = model.encode(df_templates['EventTemplate'].tolist())  # num_workers=num_workers)
        df_templates['Vector'] = [embedding.tolist() for embedding in embeddings]
        template_dict = df_templates.set_index('EventTemplate')['Vector'].to_dict()

        # 替换演化模板为最相似的旧模板
        instead_dict = find_most_similar_eventId(df_templates, df_templates_evo)
        for idx, eventId in tqdm(enumerate(df['EventId'])):
            if eventId in instead_dict:
                df.at[idx, 'EventId'] = instead_dict[eventId]


    # data preprocess
    df["Label"] = df["Label"].apply(lambda x: int(x != "-"))

    df['datetime'] = pd.to_datetime(df["Date"] + " " + df['Time'], format='%Y.%m.%d %H:%M:%S')
    df['timestamp'] = df["datetime"].values.astype(np.int64) // 10 ** 9
    df['deltaT'] = df['datetime'].diff() / np.timedelta64(1, 's')
    df['deltaT'].fillna(0)

    # sampling with sliding window
    deeplog_df = sliding_window(df[["timestamp", "Label", "EventId", "deltaT"]],
                                para={"window_size": float(window_size)*60, "step_size": float(step_size) * 60}
                                )
    output_dir += window_name

    #########
    # Train #
    #########
    df_normal = deeplog_df[deeplog_df["Label"] == 0]
    df_normal = df_normal.sample(frac=1, random_state=12).reset_index(drop=True) #shuffle
    normal_len = len(df_normal)
    train_len = int(train_ratio) if train_ratio >= 1 else int(normal_len * train_ratio)

    train = df_normal[:train_len]
    if process_evolution:
        deeplog_file_generator(os.path.join(output_dir,'train_veolution'), train, ["EventId"])
    else:
        deeplog_file_generator(os.path.join(output_dir,'train'), train, ["EventId"])
    print("training size {}".format(train_len))


    ###############
    # Test Normal #
    ###############
    test_normal = df_normal[train_len:]
    if process_evolution:
        deeplog_file_generator(os.path.join(output_dir, 'test_normal_evolution'), test_normal, ["EventId"])
    else:
        deeplog_file_generator(os.path.join(output_dir, 'test_normal'), test_normal, ["EventId"])
    print("test normal size {}".format(normal_len - train_len))


    #################
    # Test Abnormal #
    #################
    df_abnormal = deeplog_df[deeplog_df["Label"] == 1]
    if process_evolution:
        deeplog_file_generator(os.path.join(output_dir,'test_abnormal_evolution'), df_abnormal, ["EventId"])
    else:
        deeplog_file_generator(os.path.join(output_dir,'test_abnormal'), df_abnormal, ["EventId"])

    print('test abnormal size {}'.format(len(df_abnormal)))

